use std::process;
use base64::Engine;
use base64::engine::general_purpose;
use crypto::aes;
use crypto::aes::KeySize::KeySize256;
use crypto::blockmodes::PkcsPadding;
use crypto::buffer::{BufferResult, ReadBuffer, RefReadBuffer, RefWriteBuffer, WriteBuffer};
use crypto::symmetriccipher::SymmetricCipherError;

fn main() {
    println!("Hello, world!");
}

/// aes加密
fn aes256_cbc_encrypt(data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>, SymmetricCipherError> {
    let mut encryptor = aes::cbc_encryptor(KeySize256, key, iv, PkcsPadding);
    let mut buffer = [0; 4096];
    let mut write_buffer = RefWriteBuffer::new(&mut buffer);
    let mut read_buffer = RefReadBuffer::new(data);
    let mut final_result = Vec::new();
    loop {
        let result = encryptor.encrypt(&mut read_buffer, &mut write_buffer, true)?;
        final_result.extend(write_buffer.take_read_buffer().take_remaining().iter().map(|&i| i));
        match result {
            BufferResult::BufferUnderflow => break,
            _ => continue,
        }
    }
    Ok(final_result)
}

/// aes解密
fn aes256_cbc_decrypt(data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>, SymmetricCipherError> {
    let mut decryptor = aes::cbc_decryptor(KeySize256, key, iv, PkcsPadding);
    let mut buffer = [0; 4096];
    let mut write_buffer = RefWriteBuffer::new(&mut buffer);
    let mut read_buffer = RefReadBuffer::new(data);
    let mut final_result = Vec::new();
    loop {
        let result = decryptor.decrypt(&mut read_buffer, &mut write_buffer, true)?;
        final_result.extend(write_buffer.take_read_buffer().take_remaining().iter().map(|&i| i));
        match result {
            BufferResult::BufferUnderflow => break,
            _ => continue,
        }
    }
    Ok(final_result)
}


pub fn simple_aes_encrypt(input: &str, key: &str, iv: &str) -> String {
    let key = general_purpose::STANDARD.decode(key);
    if key.is_err() {
        println!("解密失败，密钥不正确");
        process::exit(1);
    }
    let key = key.unwrap();
    let iv = general_purpose::STANDARD.decode(iv);
    if iv.is_err() {
        println!("解密失败，密钥不正确");
        process::exit(1);
    }
    let iv = iv.unwrap();
    let r = aes256_cbc_encrypt(input.as_bytes(), &key, &iv);
    if r.is_err() {
        println!("解密失败，密钥不正确");
        process::exit(1);
    }
    let r = r.unwrap();
    general_purpose::STANDARD.encode(r)
}

pub fn simple_aes_decrypt(input: &str, key: &str, iv: &str) -> String {
    let input = general_purpose::STANDARD.decode(input);
    if input.is_err() {
        println!("解密失败，密钥不正确");
        process::exit(1);
    }
    let input = input.unwrap();
    let key = general_purpose::STANDARD.decode(key);
    if key.is_err() {
        println!("解密失败，密钥不正确");
        process::exit(1);
    }
    let key = key.unwrap();
    let iv = general_purpose::STANDARD.decode(iv);
    if iv.is_err() {
        println!("解密失败，密钥不正确");
        process::exit(1);
    }
    let iv = iv.unwrap();
    let result = aes256_cbc_decrypt(input.as_slice(), &key, &iv);
    if result.is_err() {
        println!("解密失败，密钥不正确");
        process::exit(1);
    }
    let result = result.unwrap();
    let result = std::str::from_utf8(result.as_slice());
    if result.is_err() {
        println!("解密失败，密钥不正确");
        process::exit(1);
    }
    let result = result.unwrap();
    result.to_string()
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand_core::{OsRng, RngCore};

    #[test]
    fn test1() {
        // 生成随机数
        let mut rng = OsRng::default();
        let mut key = [0; 32];
        let mut iv = [0; 16];
        rng.fill_bytes(&mut key);
        rng.fill_bytes(&mut iv);
        let key = general_purpose::STANDARD.encode(&key);
        println!("key:{}", key);
        let iv = general_purpose::STANDARD.encode(&iv);
        println!("iv:{}", iv);
        // 原始密码
        let password = "这是原始密码";
        println!("原始密码:{password}");
        // 加密
        let encrypted_data = simple_aes_encrypt(password, &key, &iv);
        println!("加密后的密码:{encrypted_data}");
        assert_ne!(encrypted_data, password);
        assert!(!encrypted_data.contains("这是"));
        assert!(!encrypted_data.contains("原始"));
        assert!(!encrypted_data.contains("密码"));
        // 解密
        let decrypted_data = simple_aes_decrypt(&encrypted_data, &key, &iv);
        println!("解密后的密码:{decrypted_data}");
        assert_eq!(password, decrypted_data);
    }

    #[test]
    fn test2() {
        let encrypted_data = "CJOQP2uP90/Mk3TxBR1YZgYG5Gz3GkCqHbSW/zHprog=";
        let key = "pPA8svwdYr8G8ySHYsayD3UuoKNMxOZ83uAXO4N/nO4=";
        let iv = "nR7c1bXQpEvYpjqf2ZpLBg==";
        let decrypted_data = simple_aes_decrypt(&encrypted_data, &key, &iv);
        assert_eq!("这是原始密码", decrypted_data);
    }
}